##################################################
# (1) init ==> B <= 0
# (2) unsafe ==> B > 0
# (3) domain /\ x /\ x' ==> B(x') - (1-alpha) * B(x) < 0
##################################################
# max(B(x))  s.t. x in init
# if max(B)<=0: condition (1) is OK
import math
import time
import cvxpy as cp
import dccp
import gurobipy as gp
from gurobipy import GRB

import verify_dccp.getMathExpress2 as express
from prob import *
import numpy as np
import verify_piecewise_utils
from NN import cdinn2
import torch


def verif(model, x_len):
    cex_init = []
    res_init = []
    cex_unsafe = []
    res_unsafe = []
    cex_domain = []
    res_domain = []

    x = cp.Variable(DIM)  # x_0, x_1
    v = cp.Variable()

    delta_x3 = cp.Variable()

    f1, f2 = express.getExpr(x, model)

    # print("curvature of f1:", f1.curvature)
    # print("curvature of f2:", f2.curvature)

    # print("=====================X_INIT=========================")
    objective = cp.Maximize(f1 - v)
    constraints = [v == f2]

    cons1 = [cp.square(x[0] - 1.5) + cp.square(x[1]) <= 0.25]  # x=(1.5, 0), r=0.5
    cons2 = [-1.8 <= x[0], x[0] <= -1.2, -0.1 <= x[1], x[1] <= 0.1]  # -1.8<=x0<=-1.2, -0.1<=x1<=0.1
    cons3 = [-1.4 <= x[0], x[0] <= -1.2, -0.5 <= x[1], x[1] <= 0.1]  # -1.4<=x0<=-1.2, -0.5<=x1<=0.1
    constraints123 = [cons1, cons2, cons3]

    for i in range(len(constraints123)):
        problem = cp.Problem(objective, constraints + constraints123[i])
        # print("problem is DCCP:", dccp.is_dccp(problem))
        result = problem.solve(method='dccp')
        # print("x.value:", x.value)  # [ 1.34102711 -0.47405446]
        # print("v.value:", v.value)  # 4.158081062071981
        # print("status:", problem.status) # Converged
        # print("cost value0 =", result[0])  # -1.161325675958789
        if result[0] > 0.0:
            cex_init.append(x.value)
            res_init.append(result[0])

    # print("=====================X_UNSAFE=========================")
    objective = cp.Minimize(v - f2)
    constraints = [v == f1]

    cons1 = [cp.square(x[0]+1) + cp.square(x[1]+1) <= 0.16]  # x=(-1, -1), r=0.4
    cons2 = [0.4 <= x[0], x[0] <= 0.6, 0.1 <= x[1], x[1] <= 0.5]  # 0.4<=x0<=0.6, 0.1<=x1<=0.5
    cons3 = [0.4 <= x[0], x[0] <= 0.8, 0.1 <= x[1], x[1] <= 0.3]  # 0.4<=x0<=0.8, 0.1<=x1<=0.3
    constraints123 = [cons1, cons2, cons3]

    for i in range(len(constraints123)):
        problem = cp.Problem(objective, constraints + constraints123[i])
        # print("problem is DCCP:", dccp.is_dccp(problem))
        # x.value = np.array([1, 2])
        result = problem.solve(method='dccp')
        # print("x.value:", x.value)
        # print("v.value:", v.value)
        # print("status:", problem.status)
        # print("cost value0 =", result[0])
        if result[0] < 0.0:
            cex_unsafe.append(x.value)
            res_unsafe.append(result[0])

    # print("=====================X_DOMAIN=========================")
    # fw = open('log/verify_eg4_0-01.txt', 'w+', encoding='utf-8')

    # eg2
    lb_list = [-3, 0]
    ub_list = [0, 2.5]

    lb_domain = domain_min[0]  # -3
    ub_domain = domain_max[0]  # 2.5
    # print("ub: ", ub)
    ly_domain = domain_min[1]  # -2
    uy_domain = domain_max[1]  # 1

    for i in range(len(lb_list)):
        lb = lb_list[i]
        ub = ub_list[i]
        x_l = lb
        x_r = min(round(x_l + x_len, 5), ub)
        while x_l < ub:
            # print("================[" + str(x_l) + "," + str(x_r) + "]===================")

            y_l = verify_piecewise_utils.func_x3(x_l)
            y_r = verify_piecewise_utils.func_x3(x_r)
            k, b_u, b_l = verify_piecewise_utils.linear_function_x3(x_l, x_r, y_l, y_r)

            delta_b_x3 = b_u - b_l

            x0_next = time_step * x[1] + x[0]
            x1_next = time_step * (- x[0] - x[1] + (k * x[0] + b_l + delta_x3) / 3.0) + x[1]  # set power(x[:, 0], 3) == k * x[0] + b

            f1_next, f2_next = express.getExpr([x0_next, x1_next], model)

            u = cp.Variable()
            objective = cp.Maximize((f1_next - v) - (1 - alpha) * (u - f2))
            constraints = [v == f2_next, u == f1]
            constraints += [x_l <= x[0], x[0] <= x_r]
            constraints += [ly_domain <= x[1], x[1] <= uy_domain]
            constraints += [lb_domain <= x0_next, x0_next <= ub_domain]
            constraints += [ly_domain <= x1_next, x1_next <= uy_domain]
            constraints += [0 <= delta_x3, delta_x3 <= delta_b_x3]

            problem = cp.Problem(objective, constraints)
            # print("problem is DCCP:", dccp.is_dccp(problem))
            # x.value = np.array([1, 2])
            result = problem.solve(method='dccp')
            # print("x.value:", x.value)
            # print("delta.value:", delta.value)
            # print("v.value:", v.value)
            # print("status:", problem.status)
            # print("cost value =", result[0])
            # fw.write("================[" + str(x_l) + "," + str(x_r) + "]===================")
            # fw.write('\n')
            # fw.write(str(result[0]))
            # fw.write('\n')
            # fw.write('\n')

            if result[0] > 0.0:
                cex_domain.append(x.value)
                res_domain.append(result[0])

            x_l = x_r
            x_r = min(round(x_l + x_len, 5), ub)
    # fw.close()
    # print("==============cex and res=============")
    print("cex_init: ", cex_init)
    print("res_init: ", res_init)
    print("cex_unsafe: ", cex_unsafe)
    print("res_unsafe: ", res_unsafe)
    print("cex_domain: ", cex_domain)
    print("res_domain: ", res_domain)
    return cex_init, cex_unsafe, cex_domain


if __name__ == "__main__":
    model = cdinn2.gen_nn()
    model.load_state_dict(torch.load('../cdinn-model/cdinn-trained_eg4_cdinn2_10_2_0-1_relu_update.pt'), strict=True)
    time_start_verify = time.time()
    verif(model)
    time_end_verify = time.time()
    print("verifing totally costs:", time_end_verify - time_start_verify)



